# agent/paper_tracer.py
from typing import Dict, List, Any, Tuple
import streamlit as st
from config import Config
from agent.nodes import PaperPalNodes
from prompts.trace_prompts import TRACE_ANALYSIS_PROMPT
import re
import json


class PaperTracer:
    """论文溯源分析器 - 提供深度学术谱系分析"""

    def __init__(self, api_key: str):
        self.api_key = api_key
        self.nodes = PaperPalNodes(api_key)

    def analyze_citation_impact(self, paper_title: str, paper_abstract: str) -> Tuple[str, str]:
        """执行完整的学术谱系分析（单次API调用）"""
        try:
            # 构建提示词 - 一次性获取所有内容
            prompt = TRACE_ANALYSIS_PROMPT.format(
                paper_title=paper_title,
                paper_abstract=paper_abstract[:3000]  # 摘要截断
            )

            # 调用GLM模型进行溯源分析
            analysis_result = self.nodes.call_glm(prompt, max_tokens=3500)

            # 提取Mermaid代码块
            citation_graph = self.extract_mermaid_code(analysis_result)

            # 清理图谱代码
            citation_graph = self.clean_mermaid_code(citation_graph)

            # 从完整结果中移除Mermaid代码块
            if citation_graph:
                citation_analysis = analysis_result.replace(f"```mermaid\n{citation_graph}\n```", "")
            else:
                citation_analysis = analysis_result

            return citation_analysis, citation_graph
        except Exception as e:
            st.error(f"学术谱系分析失败: {str(e)}")
            return f"分析失败: {str(e)}", ""

    def extract_mermaid_code(self, analysis_text: str) -> str:
        """从分析文本中提取Mermaid代码块"""
        if "```mermaid" in analysis_text:
            # 提取代码块内容
            mermaid_part = analysis_text.split("```mermaid")[1].split("```")[0].strip()
            return mermaid_part
        return ""

    def clean_mermaid_code(self, mermaid_text: str) -> str:
        """清理Mermaid代码，确保格式正确"""
        # 清理可能的额外字符
        mermaid_text = re.sub(r'^[`]+', '', mermaid_text)
        mermaid_text = re.sub(r'[`]+$', '', mermaid_text)
        mermaid_text = mermaid_text.strip()

        # 确保以graph开头
        if not mermaid_text.startswith("graph"):
            # 尝试找到graph开头的行
            lines = mermaid_text.split('\n')
            for i, line in enumerate(lines):
                if line.strip().startswith("graph"):
                    return '\n'.join(lines[i:])

            # 如果找不到，添加默认graph声明
            mermaid_text = "graph TD\n" + mermaid_text

        # 确保使用正确的方向
        if "graph TD" not in mermaid_text and "graph LR" not in mermaid_text:
            mermaid_text = mermaid_text.replace("graph", "graph TD", 1)

        return mermaid_text

    def analyze_citation_impact_node(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """溯源分析节点 - 集成到工作流中"""
        st.info("🔍 正在进行深度学术谱系分析...")

        # 获取论文信息
        paper_title = state.get("file_name", "目标论文")
        paper_abstract = state.get("summary", "")

        # 执行分析（单次API调用）
        citation_analysis, citation_graph = self.analyze_citation_impact(paper_title, paper_abstract)

        # 更新状态
        state.update({
            "citation_analysis": citation_analysis,
            "citation_graph": citation_graph,
            "trace_complete": True
        })

        st.success("✅ 学术谱系分析完成！")
        return state